# ==================== 实验配置区 ====================
import os

import torch

from MultiModalEmotionModel import load_bert_model, load_and_process_data, k_fold_cross_validation, \
    k_fold_cross_validation_config

EXPERIMENT_CONFIGS = [

    {  # 示例配置1：基础BERT模型
        "name": "nor_bert_1e-5",
        "bert_path": "./bert/nor_mosei_lr_1.000000e-05",
        "text_csv": "data/cmu_mosei/metadata/merged_data_clean.csv",
        "video_dir": "./data/out_feature_mosei",
        "cache_path":"data/cmu_mosei/tensor",
        "max_video_dim": 1463,
        "text_dim": 768,
        "num_classes": 7,
        "lr": 0.0001,
        "epochs": 100,
        "k_folds": 5
    },
    {  # 示例配置1：基础BERT模型
        "name": "nor_bert_2e-5",
        "bert_path": "./bert/nor_mosei_lr_2.000000e-05",
        "text_csv": "data/cmu_mosei/metadata/merged_data_clean.csv",
        "video_dir": "./data/out_feature_mosei",
        "cache_path":"data/cmu_mosei/tensor",
        "max_video_dim": 1463,
        "text_dim": 768,
        "num_classes": 7,
        "lr": 0.0001,
        "epochs": 100,
        "k_folds": 5
    },
    {  # 示例配置1：基础BERT模型
        "name": "nor_bert_3e-5",
        "bert_path": "./bert/nor_mosei_lr_3.000000e-05",
        "text_csv": "data/cmu_mosei/metadata/merged_data_clean.csv",
        "video_dir": "./data/out_feature_mosei",
        "max_video_dim": 1463,
        "text_dim": 768,
        "num_classes": 7,
        "lr": 0.0001,
        "epochs": 100,
        "k_folds": 5
    },
    {  # 示例配置1：基础BERT模型
        "name": "nor_bert_5e-5",
        "bert_path": "./bert/nor_mosei_lr_5.000000e-05",
        "text_csv": "data/cmu_mosei/metadata/merged_data_clean.csv",
        "video_dir": "./data/out_feature_mosei",
        "max_video_dim": 1463,
        "text_dim": 768,
        "num_classes": 7,
        "lr": 0.0001,
        "epochs": 100,
        "k_folds": 5
    },
    {  # small
        "name": "small_bert_1e-5",
        "bert_path": "./bert/small_mosei_lr_1.000000e-05",
        "text_csv": "data/cmu_mosei/metadata/merged_data_clean_small.csv",
        "video_dir": "./data/out_feature_mosei",
        "max_video_dim": 1463,
        "text_dim": 768,
        "num_classes": 7,
        "lr": 0.0001,
        "epochs": 100,
        "k_folds": 5
    },
    {  # BERT模型
        "name": "small_bert_2e-5",
        "bert_path": "./bert/small_mosei_lr_2.000000e-05",
        "text_csv": "data/cmu_mosei/metadata/merged_data_clean_small.csv",
        "video_dir": "./data/out_feature_mosei",
        "max_video_dim": 1463,
        "text_dim": 768,
        "num_classes": 7,
        "lr": 0.0001,
        "epochs": 100,
        "k_folds": 5
    },    {  # BERT模型
        "name": "small_bert_3e-5",
        "bert_path": "./bert/small_mosei_lr_3.000000e-05",
        "text_csv": "data/cmu_mosei/metadata/merged_data_clean_small.csv",
        "video_dir": "./data/out_feature_mosei",
        "max_video_dim": 1463,
        "text_dim": 768,
        "num_classes": 7,
        "lr": 0.0001,
        "epochs": 100,
        "k_folds": 5
    },{
        "name": "small_bert_5e-5",
        "bert_path": "./bert/small_mosei_lr_5.000000e-05",
        "text_csv": "data/cmu_mosei/metadata/merged_data_clean_small.csv",
        "video_dir": "./data/out_feature_mosei",
        "cache_path":"data/cmu_mosei/tensor",
        "max_video_dim": 1463,
        "text_dim": 768,
        "num_classes": 7,
        "lr": 0.001,
        "epochs": 100,
        "k_folds": 5
     },

        {  # max模型
            "name": "max_bert_1e-5",
            "bert_path": "./bert/max_mosei_lr_1.000000e-05",
            "text_csv": "data/cmu_mosei/metadata/merged_data_clean_max.csv",
            "video_dir": "./data/out_feature_mosei",
            "max_video_dim": 1463,
            "text_dim": 768,
            "num_classes": 6,
            "lr": 0.0001,
            "epochs": 100,
            "k_folds": 5
        }, {  # BERT模型
            "name": "max_bert_2e-5",
            "bert_path": "./bert/max_mosei_lr_2.000000e-05",
            "text_csv": "data/cmu_mosei/metadata/merged_data_clean_max.csv",
            "video_dir": "./data/out_feature_mosei",
            "max_video_dim": 1463,
            "text_dim": 768,
            "num_classes": 6,
            "lr": 0.0001,
            "epochs": 100,
            "k_folds": 5
        },
    {  # BERT模型
        "name": "max_bert_3e-5",
        "bert_path": "./bert/max_mosei_lr_3.000000e-05",
        "text_csv": "data/cmu_mosei/metadata/merged_data_clean_max.csv",
        "video_dir": "./data/out_feature_mosei",
        "max_video_dim": 1463,
        "text_dim": 768,
        "num_classes": 6,
        "lr": 0.0001,
        "epochs": 100,
        "k_folds": 5
    }, {  # BERT模型
        "name": "max_bert_5e-5",
        "bert_path": "./bert/max_mosei_lr_5.000000e-05",
        "text_csv": "data/cmu_mosei/metadata/merged_data_clean_max.csv",
        "video_dir": "./data/out_feature_mosei",
        "max_video_dim": 1463,
        "text_dim": 768,
        "num_classes": 6,
        "lr": 0.0001,
        "epochs": 100,
        "k_folds": 5
    },{  # mer
        "name": "mer_bert_1e-5",
        "bert_path": "./bert/mer_lr_1.00e-05",
        "text_csv": "data/mer/text/merged_file_with_labels.csv",
        "cache_path": "data/mer/tensor",
        "video_dir": "./data/mer/out_feature",
        "max_video_dim": 1463,
        "text_dim": 768,
        "num_classes": 6,
        "lr": 0.0001,
        "epochs": 100,
        "k_folds": 5
    },
    {  # BERT模型
        "name": "mer_bert_2e-5",
        "bert_path": "./bert/mer_lr_2.00e-05",
        "text_csv": "data/mer/text/merged_file_with_labels.csv",
        "cache_path": "data/mer/tensor",
        "video_dir": "./data/mer/out_feature",
        "max_video_dim": 1463,
        "text_dim": 768,
        "num_classes": 6,
        "lr": 0.0001,
        "epochs": 100,
        "k_folds": 5
    }, {  # BERT模型
        "name": "mer_bert_3e-5",
        "bert_path": "./bert/mer_lr_3.00e-05",
        "text_csv": "data/mer/text/merged_file_with_labels.csv",
        "cache_path": "data/mer/tensor",
        "video_dir": "./data/mer/out_feature",
        "max_video_dim": 1463,
        "text_dim": 768,
        "num_classes": 6,
        "lr": 0.0001,
        "epochs": 100,
        "k_folds": 5
    }, {
        "name": "mer_bert_5e-5",
        "bert_path": "./bert/mer_lr_5.00e-05",
        "text_csv": "data/mer/text/merged_file_with_labels.csv",
        "cache_path": "data/mer/tensor",
        "video_dir": "./data/mer/out_feature",
        "max_video_dim": 1463,
        "text_dim": 768,
        "num_classes": 6,
        "lr": 0.001,
        "epochs": 100,
        "k_folds": 5
    }
]


# ==================== 主执行函数 ====================
def main():
    # 创建总结果目录
    os.makedirs("results", exist_ok=True)


    # 遍历所有实验配置
    for cfg in EXPERIMENT_CONFIGS :
        print(f"\n{'=' * 40}")
        print(f" 开始实验: {cfg['name']} ")
        print(f"{'=' * 40}")

        # 加载BERT模型
        tokenizer, bert_model = load_bert_model(cfg["bert_path"])

        # 处理数据
        video_feats, text_feats, labels = load_and_process_data(
            text_csv_path=cfg["text_csv"],
            video_features_dir=cfg["video_dir"],
            cache_path=cfg["cache_path"],
            tokenizer=tokenizer,
            bert_model=bert_model,
            max_video_feature_dim=cfg["max_video_dim"],
            name = cfg['name']
        )

        # 执行K折验证
        k_fold_cross_validation_config(
            video_features=video_feats,
            text_features=text_feats,
            labels=labels,
            num_classes=cfg["num_classes"],
            k=cfg["k_folds"],
            lr=cfg["lr"],
            epochs=cfg["epochs"],
            model_name=cfg["name"]  # 传递模型名称参数
        )


if __name__ == "__main__":
    main()